import numpy as np

from .dmp import DMP
from ..canonical_system import CanonicalSystem, DiscreteCanonicalSystem


class DiscreteDMP(DMP):

    def __init__(self, n_bfs=100, alpha_y=60.0, beta_y=15.0, alpha_x=5.0) -> None:
        super().__init__()

        self.alpha_y = alpha_y
        self.beta_y = beta_y

        self.cs: CanonicalSystem = DiscreteCanonicalSystem(alpha=alpha_x)

        self.timesteps = 0
        self.n_dmps = 0
        self.dt = 0.0
        self.y0 = np.array([])
        self.goal = np.array([])

        self.n_bfs = n_bfs
        self.w = np.array([])
        self.psi_centers = np.zeros(self.n_bfs)
        self._calculate_phi_centers()
        self.psi_h = np.zeros(self.n_bfs)
        self._calculate_phi_h()

        self._reset()

    def _reset(self):
        self.y = self.y0.copy()
        self.dy = np.zeros_like(self.y)
        self.ddy = np.zeros_like(self.y)
        self.cs.reset()

    def _calculate_phi_centers(self):
        t_centers = np.linspace(0, self.cs.run_time, self.n_bfs)

        for i in range(t_centers.size):
            self.psi_centers[i] = self.cs.interpolate(t_centers[i])

    def _calculate_phi_h(self):
        self.psi_h = np.ones(self.n_bfs) * self.n_bfs ** 1.5 / self.psi_centers / self.cs.alpha

    def _calculate_psi(self, x: [float, np.ndarray]) -> [float, np.ndarray]:
        if isinstance(x, np.ndarray):
            x = x.reshape((-1, 1))

        psi = np.exp(-self.psi_h * (x - self.psi_centers) ** 2)
        return psi

    def _calculate_weights(self, f_target):
        x_track = self.cs.run()
        psi_track = self._calculate_psi(x_track)

        for d in range(self.n_dmps):
            for b in range(self.n_bfs):
                numer = np.sum(x_track * psi_track[:, b] * f_target[:, d])
                denom = np.sum(x_track ** 2 * psi_track[:, b])

                self.w[d, b] = numer / denom

        self.w = np.nan_to_num(self.w)

    def learning(self, y_demo):
        if y_demo.ndim == 1:
            y_demo = y_demo.reshape((-1, 1))

        self.y0 = y_demo[0, :].copy()
        self.goal = y_demo[-1, :].copy()

        self.timesteps = y_demo.shape[0]
        self.n_dmps = y_demo.shape[1]
        self.dt = self.cs.run_time / (self.timesteps - 1)
        self.cs.dt = self.dt

        self.w = np.zeros((self.n_dmps, self.n_bfs))

        dy_demo = np.gradient(y_demo, axis=0) / self.dt
        ddy_demo = np.gradient(dy_demo, axis=0) / self.dt

        x_track = self.cs.run()
        f_target = np.zeros((self.timesteps, self.n_dmps))
        for d in range(self.n_dmps):
            k = self.alpha_y
            f_target[:, d] = (ddy_demo[:, d] - self.alpha_y * (
                    self.beta_y * (self.goal[d] - y_demo[:, d]) - dy_demo[:, d])) / k + x_track * (
                                     self.goal[d] - self.y0[d])

        self._calculate_weights(f_target)

        self._reset()

    def reproduce(self, tau=1.0, initial=None, goal=None):
        timesteps = round(self.timesteps / tau)

        if initial is not None:
            self.y0 = initial.copy()

        if goal is not None:
            self.goal = goal.copy()

        self._reset()

        y_reproduce = np.zeros((timesteps, self.n_dmps))
        dy_reproduce = np.zeros((timesteps, self.n_dmps))
        ddy_reproduce = np.zeros((timesteps, self.n_dmps))

        for t in range(timesteps):
            y_reproduce[t], dy_reproduce[t], ddy_reproduce[t] = self.step(tau=tau)

        return y_reproduce

    def step(self, tau=1.0):
        x = self.cs.step(tau)

        psi = self._calculate_psi(x)

        for d in range(self.n_dmps):
            k = self.alpha_y
            f = k * (np.dot(psi, self.w[d]) * x / np.sum(psi)) - k * (self.goal[d] - self.y0[d]) * x

            self.y[d] += tau * self.dy[d] * self.dt
            self.ddy[d] = self.alpha_y * (self.beta_y * (self.goal[d] - self.y[d]) - self.dy[d]) + f
            self.dy[d] += tau * self.ddy[d] * self.dt

        return self.y, self.dy, self.ddy
